'''
Created on Oct 16, 2009

@author: mkiyer
'''

import os
import sys
import collections
from optparse import OptionParser
from bx.intervals.intersection import Interval, IntervalTree

__path_to_tables = '/lab/mkiyer/refdb/ucsc_tables'

def _parse_targetscan_sites():
    '''
    #bin    chrom   chromStart      chromEnd        name    score   strand
    '''
    # read list of mirnas
    fhd = open(os.path.join(__path_to_tables, 'targetScanS.txt'))
    for thisline in fhd:
        thisline = thisline.strip()
        if thisline == None:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        mirna = type('TargetSite', (object,), dict())
        mirna.chrom = thisfields[1]
        mirna.start = int(thisfields[2])
        mirna.end = int(thisfields[3])
        namefields = thisfields[4].split(':')
        mirna.gene = namefields[0]
        mirna.name = ''.join(namefields[1])
        mirna.score = int(thisfields[5])
        mirna.strand = thisfields[6]
        yield mirna


def build_mirna_interval_trees():
    # build interval tree of mirna target sites
    mirna_trees = collections.defaultdict(lambda: IntervalTree())  
    for ts in _parse_targetscan_sites():
        mirna_trees[ts.chrom].insert_interval(Interval(ts.start, ts.end, value=ts))
    return mirna_trees

__target_sites = None

def get_mirna_target_sites(chrom, start, end):
    global __target_sites
    if not __target_sites:
        __target_sites = build_mirna_interval_trees()
    hits = __target_sites[chrom].find(start, end)
    return [hit.value.name for hit in hits]
    

if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options] <file>")
    optionparser.add_option("--chrom", type="int", dest="chrom",
                            default=2)
    optionparser.add_option("--strand", type="int", dest="strand",
                            default=3)
    optionparser.add_option("--start", type="int", dest="start",
                            default=4)
    optionparser.add_option("--end", type="int", dest="end",
                            default=5)
    optionparser.add_option("--polya", type="int", dest="polya",
                            default=6)    
    (options, args) = optionparser.parse_args()
    maxcol = max(options.chrom, options.strand, options.start, options.end, options.polya)
    
    for thisline in open(args[0]):
        thisline = thisline.strip()
        if not thisline:
            continue
        if thisline.startswith('#'):
            continue
        thisfields = thisline.split('\t')
        if len(thisfields) < maxcol+1:
            # TODO: print error message
            continue
        chrom = thisfields[options.chrom]
        strand = thisfields[options.strand]
        start = int(thisfields[options.start])
        end = int(thisfields[options.end])
        polya = int(thisfields[options.polya])
        if strand == '+':
            interval = (polya, end)
        else:
            interval = (start, polya)
        
        target_sites = get_mirna_target_sites(chrom, interval[0], interval[1])
        if len(target_sites) == 0:
            target_sites = ['None']
        newline = '\t'.join([thisline, ','.join(target_sites)])
        sys.stdout.write('%s\n' % newline)
        